import torch
import torch.nn as nn


class GroupNorm(nn.Module):
    def __init__(self, c_num, group_num=16, eps=1e-10):
        super(GroupNorm, self).__init__()
        self.group_num = group_num
        self.gamma = nn.Parameter(torch.ones(c_num, 1, 1))
        self.beta = nn.Parameter(torch.zeros(c_num, 1, 1))
        self.eps = eps

    def forward(self, x):
        N, C, H, W = x.size()

        x = x.view(N, self.group_num, -1)

        mean = x.mean(dim=2, keepdim=True)
        std = x.std(dim=2, keepdim=True)

        x = (x - mean) / (std + self.eps)
        x = x.view(N, C, H, W)

        return x * self.gamma + self.beta

if __name__ == '__main__':
    # gn = GroupNorm(32)
    # x = torch.randn((3,32,128,128))
    # output = gn(x)
    # print(output.shape)
    x = torch.ones((3,2,4,4))
    e = torch.ones((2,1,1)) * 2
    output = x * e
    print(output.shape)